package com.nc.mem.serial.graph.container;

import gnu.trove.impl.Constants;
import gnu.trove.map.hash.TIntObjectHashMap;

import java.io.PrintWriter;
import java.lang.reflect.Modifier;
import java.util.Collection;
import java.util.Comparator;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.SortedMap;
import java.util.SortedSet;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.util.ASMifier;
import org.objectweb.asm.util.TraceClassVisitor;

import com.nc.mem.GraphClassLoader;
import com.nc.mem.serial.graph.Instantiator;
import com.nc.mem.serial.graph.Representation;
import com.nc.mem.serial.graph.boot.ProvisionService;
import com.nc.mem.serial.graph.boot.ProvisionServiceFactory;
import com.nc.mem.serial.graph.boot.RegistrationListener;


public final class CollectionRegistry implements RegistrationListener, Opcodes {
	public static boolean TRACE = false;

	final TIntObjectHashMap<Instantiator<Collection<Object>>> COL_INSTANTIATOR;

	final TIntObjectHashMap<Instantiator<Map<Object, Object>>> MAP_INSTANTIATOR;

	final IdentityHashMap<Class<?>, Class<?>> REPLACEMENTS;

	public static final CollectionRegistry INSTANCE = new CollectionRegistry();

	private CollectionRegistry() {
		this(ProvisionServiceFactory.resolveOrFallback(null));
	}

	private CollectionRegistry(final Iterable<ProvisionService> services) {
		COL_INSTANTIATOR = new TIntObjectHashMap<Instantiator<Collection<Object>>>(4, Constants.DEFAULT_LOAD_FACTOR, -1);
		MAP_INSTANTIATOR = new TIntObjectHashMap<Instantiator<Map<Object, Object>>>(4, Constants.DEFAULT_LOAD_FACTOR, -1);
		REPLACEMENTS = new IdentityHashMap<Class<?>, Class<?>>();

		init(services);
	}

	private void init(final Iterable<ProvisionService> services) {

		for (final ProvisionService service : services) {
			final Iterable<Entry<Class<?>, Class<?>>> toReplace = service.installableReplacements();

			if (toReplace != null) {
				for (final Entry<Class<?>, Class<?>> e : toReplace) {
					final Class<?> key = e.getKey();
					final Class<?> value = e.getValue();

					if (Collection.class.isAssignableFrom(key) && Collection.class.isAssignableFrom(value) || //
							Map.class.isAssignableFrom(key) && Map.class.isAssignableFrom(value)) {
						REPLACEMENTS.put(key, value);
					}
				}
			}
		}
	}

	public Collection<Object> newCollection(final int id, final Comparator<?> c) {
		return COL_INSTANTIATOR.get(id).instantiate(c);
	}

	public Collection<Object> newCollection(final int id, final int sz) {
		return COL_INSTANTIATOR.get(id).instantiate(sz);
	}

	public Map<Object, Object> newMap(final int id, final Comparator<?> c) {
		return MAP_INSTANTIATOR.get(id).instantiate(c);
	}

	public Map<Object, Object> newMap(final int id, final int sz) {
		return MAP_INSTANTIATOR.get(id).instantiate(sz);
	}

	@SuppressWarnings("unchecked")
	@Override
	public synchronized void onClassRegistered(final Class<?> clazz, final int id) {
		if (!clazz.isInterface() && !Modifier.isAbstract(clazz.getModifiers())) {
			@SuppressWarnings("rawtypes")
			final TIntObjectHashMap target = Map.class.isAssignableFrom(clazz) ? MAP_INSTANTIATOR : Collection.class.isAssignableFrom(clazz) ? COL_INSTANTIATOR : null;

			if (target != null) {
				Object instantiator = target.get(id);

				if (instantiator == null) {
					final Class<?> rep = REPLACEMENTS.get(clazz);

					if (rep != null) {
						final int repId = Representation.classToId(rep);

						instantiator = target.get(repId);

						if (instantiator == null) {
							instantiator = weaveInstantiator(rep);
							target.put(repId, instantiator);
						}
					} else {
						instantiator = weaveInstantiator(clazz);
					}

					target.put(id, instantiator);
				}
			}
		}
	}

	private Instantiator<?> weaveInstantiator(final Class<?> clazz) {

		final Type type = Type.getType(Instantiator.class);
		final String clazzIn = Type.getInternalName(clazz);
		final String in = type.getInternalName();
		final String desc = type.getDescriptor();

		final Type oType = Type.getType(Object.class);
		final String oDesc = oType.getDescriptor();

		final boolean isMap = Map.class.isAssignableFrom(clazz);
		final Type cType = isMap ? Type.getType(Map.class) : Type.getType(Collection.class);
		final String cDesc = cType.getDescriptor();

		final String sigPre = oDesc + desc.substring(0, desc.length() - 1) + "<";
		String genColDesc = cDesc.substring(0, cDesc.length() - 1) + "<" + oDesc;
		genColDesc = isMap ? genColDesc + oDesc + ">;" : genColDesc + ">;";
		final String classSig = sigPre + genColDesc + ";>";

		final String targetName = in + "Of" + clazz.getSimpleName();

		final ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
		MethodVisitor mv;

		boolean szAvailable = false;
		boolean compAvailable = false;
		try {
			if (Modifier.isPublic(clazz.getDeclaredConstructor(int.class).getModifiers())) {
				szAvailable = true;
			}
		} catch (final Exception e) {
			try {
				clazz.getDeclaredConstructor();
			} catch (final Exception ex) {
				throw new RuntimeException(ex);
			}
		}

		if (SortedMap.class.isAssignableFrom(clazz) || SortedSet.class.isAssignableFrom(clazz)) {
			try {
				if (Modifier.isPublic(clazz.getDeclaredConstructor(Comparator.class).getModifiers())) {
					compAvailable = true;
				}
			} catch (final Exception e) {
				try {
					clazz.getDeclaredConstructor();
				} catch (final Exception ex) {
					throw new RuntimeException(ex);
				}
			}
		}

		// definition
		cw.visit(V1_6, ACC_PUBLIC + ACC_SUPER, targetName, classSig, oType.getInternalName(), new String[]{ Type.getInternalName(Instantiator.class) });

		// ctor
		mv = cw.visitMethod(ACC_PUBLIC, "<init>", "()V", null, null);
		mv.visitCode();
		mv.visitVarInsn(ALOAD, 0);
		mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V");
		mv.visitInsn(RETURN);
		mv.visitMaxs(1, 1);
		mv.visitEnd();

		// impl - size
		mv = cw.visitMethod(ACC_PUBLIC, "instantiate", "(I)" + cDesc, "(I)" + genColDesc, null);
		mv.visitCode();
		mv.visitTypeInsn(NEW, clazzIn);
		mv.visitInsn(DUP);
		if (szAvailable) {
			mv.visitVarInsn(ILOAD, 1);
			mv.visitMethodInsn(INVOKESPECIAL, clazzIn, "<init>", "(I)V");
		} else {
			mv.visitMethodInsn(INVOKESPECIAL, clazzIn, "<init>", "()V");
		}

		mv.visitInsn(ARETURN);
		mv.visitMaxs(0, 0);
		mv.visitEnd();

		// impl - comp
		final String compDesc = "(" + Type.getDescriptor(Comparator.class) + ")";
		final String compDescWC = "(L" + Type.getInternalName(Comparator.class) + "<*>;)";

		mv = cw.visitMethod(ACC_PUBLIC, "instantiate", compDesc + cDesc, compDescWC + genColDesc, null);
		mv.visitCode();
		mv.visitTypeInsn(NEW, clazzIn);
		mv.visitInsn(DUP);
		if (compAvailable) {
			mv.visitVarInsn(ALOAD, 1);
			mv.visitMethodInsn(INVOKESPECIAL, clazzIn, "<init>", compDesc + "V");
		} else {
			mv.visitMethodInsn(INVOKESPECIAL, clazzIn, "<init>", "()V");
		}

		mv.visitInsn(ARETURN);
		mv.visitMaxs(0, 0);
		mv.visitEnd();

		// synthetic bridge - size
		mv = cw.visitMethod(ACC_PUBLIC + ACC_BRIDGE + ACC_SYNTHETIC, "instantiate", "(I)Ljava/lang/Object;", null, null);
		mv.visitCode();
		mv.visitVarInsn(ALOAD, 0);
		mv.visitVarInsn(ILOAD, 1);
		mv.visitMethodInsn(INVOKEVIRTUAL, targetName, "instantiate", "(I)" + cDesc);
		mv.visitInsn(ARETURN);
		mv.visitMaxs(0, 0);
		mv.visitEnd();

		// synthetic bridge - comp
		mv = cw.visitMethod(ACC_PUBLIC + ACC_BRIDGE + ACC_SYNTHETIC, "instantiate", compDesc + "Ljava/lang/Object;", null, null);
		mv.visitCode();
		mv.visitVarInsn(ALOAD, 0);
		mv.visitVarInsn(ALOAD, 1);
		mv.visitMethodInsn(INVOKEVIRTUAL, targetName, "instantiate", compDesc + cDesc);
		mv.visitInsn(ARETURN);
		mv.visitMaxs(0, 0);
		mv.visitEnd();

		final byte[] buffer = cw.toByteArray();

		if (TRACE) {
			final PrintWriter pw = new PrintWriter(System.out);

			final ClassReader dumper = new ClassReader(buffer.clone());
			dumper.accept(new TraceClassVisitor(null, new ASMifier(), pw), 0);
		}

		final Class<?> weaved = GraphClassLoader.LOADER.defineClass(buffer);

		try {
			return (Instantiator<?>) weaved.newInstance();
		} catch (final Exception e) {
			throw new RuntimeException(e);
		}
	}
}